from typing import Any, Dict, List
from enum import Enum
import json

from loguru import logger
from nicegui import ui, app
from tasks.model.task import TaskStatus
from tasks.server_helper import ServerHelper  # Import TaskStatus from our model


def task_grid_ui(row_data: List[Dict[str, Any]], on_click=None):
    for v in row_data:
        v["progress"] = "|".join([f"{k}:{v}" for k, v in v.get("progress", {}).items()])
        v["__DEL__"] = "D"
        v["__DEL_RES__"] = "R"
    ui.aggrid(
        {
            "columnDefs": [
                {"headerName": "Name", "field": "name", "sortable": True},
                {"headerName": "Task", "field": "task", "sortable": True},
                {"headerName": "ID", "field": "task_hash", "sortable": True},
                {"headerName": "Progress", "field": "progress", "sortable": False},
                {
                    "headerName": "S",
                    "field": "status",
                    "cellClassRules": {
                        "bg-red-300": f"x == '{TaskStatus.FAILED.value}'",
                        "bg-green-300": f"x == '{TaskStatus.COMPLETED.value}'",
                        "bg-blue-300": f"x == '{TaskStatus.RUNNING.value}'",
                    },
                },
                {"headerName": "Host", "field": "site", "sortable": True},
                {
                    "headerName": "Priority",
                    "field": "priority",
                    "sortable": True,
                    "width": 30,
                },
                {
                    "headerName": "DEL",
                    "field": "__DEL__",
                    "sortable": False,
                    "width": 30,
                },
                {
                    "headerName": "Reset",
                    "field": "__DEL_RES__",
                    "sortable": False,
                    "width": 30,
                },
            ],
            "rowData": row_data,
            "domLayout": "autoHeight",
        },
        html_columns=list(range(3)),
    ).on("cellClicked", on_click)


def ui_tasks():
    # Get server_helper instance from the app context
    server_helper: ServerHelper = app.server_helper

    @ui.refreshable
    def ui_tasks_():
        def on_click(event):
            task_id = event.args["data"]["task_hash"]
            column_id = event.args["colId"]
            if column_id == "task_hash":
                ui.navigate.to(f"/task_detail/{task_id}")
            elif column_id == "task_name":
                app.storage.user["selected_tasks"] = task_id
                ui_tasks_.refresh()
            elif column_id == "__DEL__":
                logger.info(f"deleting task {task_id}")
                server_helper.delete_task(task_id)
                ui_tasks_.refresh()
            elif column_id == "__DEL_RES__":
                logger.info(f"resetting task {task_id}")
                server_helper.reset_task(task_id)
                ui_tasks_.refresh()
            elif column_id == "progress":
                # copy task_hash to clipboard
                ui.clipboard.write(task_id)
                ui.notify(f"copied to clipboard {task_id}")

        # Use server_helper to get tasks
        task_specs = server_helper.list_task_specs(offset=0, limit=50)
        row_data = []

        # Convert task specs to a format suitable for the grid
        for task_spec in task_specs:
            task_data = task_spec.model_dump()
            # Add empty progress dictionary if it doesn't exist
            if "progress" not in task_data:
                task_data["progress"] = {}
            # Add site information if available
            task_data["site"] = getattr(task_spec, "site", "")
            row_data.append(task_data)

        ui.textarea('json', value=json.dumps(row_data, indent=2)).classes("w-full h-full")

        task_grid_ui(row_data, on_click)

    timer = ui.timer(5, ui_tasks_.refresh, active=False)
    with ui.row():

        def on_change(e):
            if e.value:
                timer.activate()
                ui.notify("auto refresh is enabled")
            else:
                timer.deactivate()
                ui.notify("auto refresh is disabled")

        ui.checkbox(text="auto refresh", on_change=on_change)
    ui_tasks_()


page_config = {
    "nav_entry": [
        {
            "title": "TASKS",
            "url": "/tasks",
            "function": ui_tasks,
        }
    ]
}
